# -*- coding: utf-8 -*-
"""
core/代码预处理/vectorizer.py

代码向量化模块。
加载解析后的代码片段，使用文本嵌入模型将其转换为向量，并构建FAISS索引。
"""

import os
import json
import numpy as np
import faiss
from core.utils import logger, read_json_file, write_json_file

# --- 模拟文本嵌入模型配置 ---
# 在实际应用中，这里会使用一个真正的句子转换器模型或调用LLM的嵌入API
SIMULATED_EMBEDDING_DIM = 128 # 假设嵌入向量的维度

def get_simulated_text_embedding(text_content):
    """
    模拟获取文本内容的嵌入向量。
    为了可运行性，这里简单地基于文本长度和哈希生成一个固定维度的伪随机向量。
    
    参数:
        text_content (str): 需要嵌入的文本内容。

    返回:
        numpy.ndarray: 模拟的嵌入向量。
    """
    if not text_content:
        return np.zeros(SIMULATED_EMBEDDING_DIM, dtype=np.float32)
    
    # 使用文本的哈希和长度来生成可重复的伪随机向量
    seed = hash(text_content) % (2**32 -1) # 哈希值作为种子
    np.random.seed(seed)
    vector = np.random.rand(SIMULATED_EMBEDDING_DIM).astype(np.float32)
    # 简单地根据文本长度调整向量的范数，这只是一个非常粗糙的模拟
    norm_factor = min(len(text_content) / 100.0, 1.0) # 假设文本长度100时达到最大影响
    vector = vector * norm_factor
    return vector

def vectorize_code_fragments(parsed_fragments_file, faiss_index_path, metadata_path):
    """
    向量化代码片段并构建FAISS索引。

    参数:
        parsed_fragments_file (str): 解析后的代码片段JSON文件路径。
        faiss_index_path (str): FAISS索引的输出路径。
        metadata_path (str): 与FAISS索引对应的元数据JSON文件路径。

    返回:
        bool: 如果成功则返回True，否则返回False。
    """
    logger.info(f"开始向量化代码片段，输入文件: {parsed_fragments_file}")

    code_fragments = read_json_file(parsed_fragments_file)
    if not code_fragments:
        logger.error(f"未能加载代码片段数据从: {parsed_fragments_file}，或文件为空。")
        # 确保即使失败也创建空的输出文件，以便后续流程不因文件不存在而中断
        faiss_dir = os.path.dirname(faiss_index_path)
        meta_dir = os.path.dirname(metadata_path)
        os.makedirs(faiss_dir, exist_ok=True)
        os.makedirs(meta_dir, exist_ok=True)
        # 创建一个空的FAISS索引
        try:
            empty_index = faiss.IndexFlatL2(SIMULATED_EMBEDDING_DIM)
            faiss.write_index(empty_index, faiss_index_path)
            logger.info(f"创建了一个空的FAISS索引: {faiss_index_path}")
        except Exception as e_faiss:
            logger.error(f"创建空FAISS索引失败: {e_faiss}")
        # 创建一个空的元数据文件
        write_json_file(metadata_path, [])
        logger.info(f"创建了一个空的元数据文件: {metadata_path}")
        return False

    embeddings = []
    metadata_list = [] # 用于存储与向量顺序一致的元数据

    logger.info(f"共加载 {len(code_fragments)} 个代码片段进行向量化。")

    for i, fragment in enumerate(code_fragments):
        # 将代码和注释合并作为嵌入的文本内容
        text_to_embed = f"文件名: {os.path.basename(fragment.get("file_path", ""))}\n"
        text_to_embed += f"类型: {fragment.get("type", "")}\n名称: {fragment.get("name", "")}\n"
        text_to_embed += f"注释: {fragment.get("comment", "")}\n代码:\n{fragment.get("code", "")}"
        
        embedding_vector = get_simulated_text_embedding(text_to_embed)
        embeddings.append(embedding_vector)
        
        # 存储元数据，这里我们存储片段的索引，以便之后可以从原始列表中找回完整信息
        # 或者直接存储部分关键元数据
        metadata_list.append({
            "original_index": i, # 指向原始code_fragments列表中的索引
            "file_path": fragment.get("file_path", ""),
            "name": fragment.get("name", ""),
            "type": fragment.get("type", ""),
            "start_line": fragment.get("start_line", -1),
            "end_line": fragment.get("end_line", -1)
        })
        if (i + 1) % 100 == 0:
            logger.info(f"已处理 {i + 1}/{len(code_fragments)} 个代码片段的向量化...")

    if not embeddings:
        logger.warning("没有生成任何嵌入向量。可能是所有代码片段都为空。")
        # 同样，创建空索引和元数据
        try:
            empty_index = faiss.IndexFlatL2(SIMULATED_EMBEDDING_DIM)
            faiss.write_index(empty_index, faiss_index_path)
        except Exception as e_faiss:
             logger.error(f"创建空FAISS索引失败: {e_faiss}")
        write_json_file(metadata_path, [])
        return False

    embeddings_np = np.array(embeddings).astype("float32")
    logger.info(f"嵌入向量已生成，形状: {embeddings_np.shape}")

    # 构建FAISS索引
    try:
        index = faiss.IndexFlatL2(SIMULATED_EMBEDDING_DIM) # 使用L2距离
        # 或者使用 IndexFlatIP (内积) 如果嵌入是归一化的
        index.add(embeddings_np) # 添加向量到索引
        logger.info(f"FAISS索引已构建，包含 {index.ntotal} 个向量。")
    except Exception as e:
        logger.error(f"构建FAISS索引失败: {e}")
        return False

    # 保存FAISS索引和元数据
    faiss_dir = os.path.dirname(faiss_index_path)
    meta_dir = os.path.dirname(metadata_path)
    os.makedirs(faiss_dir, exist_ok=True)
    os.makedirs(meta_dir, exist_ok=True)

    try:
        faiss.write_index(index, faiss_index_path)
        logger.info(f"FAISS索引已成功保存到: {faiss_index_path}")
    except Exception as e:
        logger.error(f"保存FAISS索引失败: {e}")
        return False

    if write_json_file(metadata_path, metadata_list):
        logger.info(f"代码片段元数据已成功保存到: {metadata_path}")
    else:
        logger.error(f"保存代码片段元数据失败: {metadata_path}")
        # 即使元数据保存失败，索引可能已保存，但这会导致不一致
        return False
        
    return True

if __name__ == "__main__":
    project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))

    # 输入的已解析代码片段文件 (由preprocessor.py生成)
    # 为测试，我们假设preprocessor.py已经运行并生成了此文件
    # 或者，我们可以在这里创建一个模拟的输入文件
    mock_parsed_fragments_file = os.path.join(project_root, "data", "output", "parsed_code_fragments.json")
    
    # 输出的FAISS索引和元数据文件路径
    faiss_output_dir = os.path.join(project_root, "data", "output", "faiss_index")
    output_faiss_index_file = os.path.join(faiss_output_dir, "codebase.index")
    output_metadata_file = os.path.join(faiss_output_dir, "metadata.json")

    os.makedirs(os.path.dirname(mock_parsed_fragments_file), exist_ok=True)
    os.makedirs(faiss_output_dir, exist_ok=True)

    # 创建一个模拟的 parsed_code_fragments.json 用于测试
    mock_fragments_data = [
        {
            "file_path": "src/LED.c",
            "type": "function",
            "name": "LED_Init",
            "start_line": 10,
            "end_line": 20,
            "comment": "初始化LED模块",
            "code": "void LED_Init(void) {\n  // 初始化代码\n}"
        },
        {
            "file_path": "src/key.c",
            "type": "function",
            "name": "KEY_Scan",
            "start_line": 5,
            "end_line": 15,
            "comment": "扫描按键状态",
            "code": "int KEY_Scan(void) {\n  // 扫描代码\n  return 0;\n}"
        },
        {
            "file_path": "src/main.c",
            "type": "global_variable",
            "name": "g_system_tick",
            "start_line": 3,
            "end_line": 3,
            "comment": "全局系统时钟计数器",
            "code": "volatile unsigned int g_system_tick = 0;"
        }
    ]
    if not write_json_file(mock_parsed_fragments_file, mock_fragments_data):
        logger.error("创建模拟的 parsed_code_fragments.json 失败，测试可能无法正确进行。")
    else:
        logger.info(f"已创建模拟的输入文件: {mock_parsed_fragments_file}")

    logger.info(f"代码向量化器独立测试开始，输入片段文件: {mock_parsed_fragments_file}")
    success = vectorize_code_fragments(mock_parsed_fragments_file, 
                                       output_faiss_index_file, 
                                       output_metadata_file)

    if success:
        logger.info("代码向量化器测试成功完成。")
        # 可以尝试加载索引和元数据进行验证
        try:
            loaded_index = faiss.read_index(output_faiss_index_file)
            logger.info(f"成功加载FAISS索引，包含 {loaded_index.ntotal} 个向量。")
            loaded_metadata = read_json_file(output_metadata_file)
            if loaded_metadata:
                logger.info(f"成功加载元数据，包含 {len(loaded_metadata)} 条记录。")
                assert loaded_index.ntotal == len(loaded_metadata)
        except Exception as e:
            logger.error(f"加载或验证输出文件时出错: {e}")
    else:
        logger.error("代码向量化器测试失败。")

